from .module import PREPROCESS
import cv2
import numpy as np


@PREPROCESS.register_module
class Resize:

    def __init__(self, size, random_crop):
        self.size = size
        self.random_crop = random_crop

    def process(self, img):
        if self.random_crop:
            self.size[0] += 16
            self.size[0] += 16
        img = cv2.resize(img, dsize=tuple(self.size))
        if self.random_crop:
            r_x = np.random.randint(16)
            r_y = np.random.randint(16)
            img = img[r_y:r_y+self.size[1], r_x:r_x+self.size[0], :]
        return img